import torch
from  torch.autograd import Variable

"""
PyTorch： optim

到目前为止，我们已经更新了我们模型的权重，通过手动改变需要学习的参数变量的 .data 成员。
这对于像随机梯度下降这种简单的优化算法来说不是很困难，但是在实际训练神经网络时，我们常常使用更复杂的优化算法，如 AdaGrad，RMSProp， Adam等。
PyTorch 的 optim 包抽象了优化算法的思想，并提供了常用算法的实现。
在这个例子中，我们讲使用 nn 包来重新定义我们之前的模型，但是我们讲使用 optim 包提供的 Adam 算法来优化我们的模型：
"""

# N 批的大小
# D_in 输入的数据纬度
# H 隐藏层纬度
# D_out 输出数据纬度
# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# 创建随机的张量保存输入和输出，放到变量里
# Create random Tensors to hold inputs and outputs, and wrap them in Variables.
x = Variable(torch.randn(N, D_in))
y = Variable(torch.randn(N, D_out), requires_grad=False)

# 使用nn包来定义我们的模型和损失函数
# Use the nn package to define our model and loss function
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Linear(H, D_out)
)

loss_fn = torch.nn.MSELoss(reduction='sum')

# 使用optim包来定义一个优化器更新模型的权重，这里使用Adam，optim包包含一些其它优化算法
# Adam构造器的第一个参数是告诉优化器哪些变量需要更新梯度
# Use the optim package to define an Optimizer that will update the weights of
# the model for us. Here we will use Adam; the optim package contains many other
# optimization algoriths. The first argument to the Adam constructor tells the
# optimizer which Variables it should update.
# 设置学习率
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for t in range(500):
    # 正向传播： 把x传递到模型计算预测的y值
    # Forward pass: compute predicted y by passing x to the model.
    y_pred = model(x)

    # 计算和打印损失
    # Compute and print loss
    loss = loss_fn(y_pred, y)
    print(t, loss.data)

    # 反向传播前：使用优化器清零所有将要更新梯度的变量的梯度
    # Before the backward pass, use the optimizer object to zero all of the
    # gradients for the variables it will update (which are the learnable weights
    # of the model)
    optimizer.zero_grad()

    # 反向传播：计算模型损失的梯度
    # Backward pass: compute gradient of the loss with respect to model
    # parameters
    loss.backward()

    # 调用优化器的step函数对参数进行更新
    # Calling the step function on an Optimizer makes an update to its
    # parameters
    optimizer.step()